'''

|-------------<<-----------|
|--->llm推理--->action执行 --| # 四要素
'''

from langchain.tools import BaseTool
import requests

class WeatherTool(BaseTool):
    name = "weather"
    description = "用于查询指定城市的天气情况。输入应该是城市名称。"

    def _run(self, city: str) -> str:
        api_key = '你的API密钥'
        url = f"http://api.weatherapi.com/v1/current.json?key={api_key}&q={city}"
        response = requests.get(url)
        data = response.json()

        if 'error' in data:
            return "无法获取天气信息，请检查城市名称后重试。"
        else:
            weather_info = (f"当前{data['location']['name']}的天气是{data['current']['condition']['text']}, "
                            f"温度为{data['current']['temp_c']}°C。")
            return weather_info

    async def _arun(self, city: str) -> str:
        """异步运行工具的方法"""
        raise NotImplementedError("此工具不支持异步运行")


from langchain.agents import AgentExecutor, Tool

# from langchain.chains import LLMChain
from langchain import LLMChain
from langchain.prompts import PromptTemplate

# 初始化语言模型
from lc_infer import DeepSeek_R1_Distill_Qwen_LLM
# model_path = r'/home/ps/zhangxiancai/llm_deploy/bigfiles/models/DeepSeek-R1-Distill-Qwen-7B'
model_path = r'D:\code\other\LLMs\models\DeepSeek-R1-Distill-Qwen-1.5B'
llm = DeepSeek_R1_Distill_Qwen_LLM(mode_name_or_path=model_path)
# 定义决策链，用于判断问题是否与查询天气相关
decision_prompt = PromptTemplate(
    input_variables=["input"],
    template="如果问题是关于查询某个城市的天气，则回答'需要查询天气,城市为某某某'；否则回答'无需查询天气'。\n问题：{input}"
)
decision_chain = LLMChain(llm=llm, prompt=decision_prompt)

# 初始化自定义工具
tools = [WeatherTool()]


# 定义自定义Agent逻辑
class CustomAgent:
    def __init__(self, tools, decision_chain):
        self.tools = {tool.name: tool for tool in tools}
        self.decision_chain = decision_chain

    def run(self, input_text):
        # 使用决策链判断是否需要使用工具
        decision = self.decision_chain.predict(input=input_text)

        if "需要查询天气" in decision.lower():
            city = input_text.split()[-1]  # 简单处理，假设城市名是最后一个词
            return self.tools['weather']._run(city)
        else:
            return "这个问题不需要查询天气。"


# 创建自定义代理实例
custom_agent = CustomAgent(tools, decision_chain)

# 测试自定义代理
question = "北京今天的天气怎么样？"
response = custom_agent.run(question)
print(response)